#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：academic_trend_analysis
@File    ：mongodb_utils.py
@IDE     ：PyCharm
@Author  ：iyoahs
@Date    ：2025/6/20 11:15
@Describe：MongoDB工具函数
'''
from pymongo import MongoClient
from typing import List, Dict, Optional
from src.utils.logger import setup_logger
import datetime


logger = setup_logger()


class MongoDBUtils:
    def __init__(self, mongo_config: Dict):
        self.uri = mongo_config["uri"]
        self.db_name = mongo_config["db_name"]
        self.collection_name = mongo_config.get("collection_name", "papers")

    def get_papers_by_keywords(self, keywords: List[str], limit: int = 50) -> List[Dict]:
        """根据关键词获取论文"""
        try:
            with MongoClient(self.uri) as client:
                db = client[self.db_name]
                collection = db[self.collection_name]
                query = {"$or": [{"abstract": {"$regex": kw, "$options": "i"}} for kw in keywords]}
                return list(collection.find(query).limit(limit))
        except Exception as e:
            logger.error(f"Failed to query papers: {str(e)}")
            return []

    def get_papers_by_time_range(self, start_date: Optional[datetime.datetime] = None,
                             end_date: Optional[datetime.datetime] = None,
                             date_field: str = "published") -> List[Dict]:
        """
        根据时间范围筛选论文，默认使用 'published' 字段作为筛选依据。

        :param start_date: 开始时间（包含）
        :param end_date: 结束时间（不包含）
        :param date_field: 要筛选的日期字段，可选值："published" 或 "updated"
        :return: 匹配时间范围的论文列表
        """
        try:
            with MongoClient(self.uri) as client:
                db = client[self.db_name]
                collection = db[self.collection_name]

                # 构建查询条件
                query = {}
                if start_date or end_date:
                    query[date_field] = {}
                    if start_date:
                        query[date_field]["$gte"] = start_date.isoformat()
                    if end_date:
                        query[date_field]["$lt"] = end_date.isoformat()

                return list(collection.find(query))
        except Exception as e:
            logger.error(f"Failed to query papers by time range: {str(e)}")
            return []

    def get_all_processed_papers(self) -> List[Dict]:
        """获取所有处理过的论文"""
        try:
            with MongoClient(self.uri) as client:
                db = client[self.db_name]
                return list(db["processed_papers"].find())
        except Exception as e:
            logger.error(f"Failed to get processed papers: {str(e)}")
            return []

    def save_analysis_result(self, result_type: str, data: Dict) -> bool:
        """保存分析结果"""
        try:
            with MongoClient(self.uri) as client:
                db = client[self.db_name]
                collection = db["analysis_results"]
                data["type"] = result_type
                data["created_at"] = datetime.datetime.utcnow().isoformat()
                collection.insert_one(data)
                return True
        except Exception as e:
            logger.error(f"Failed to save analysis result: {str(e)}")
            return False

    def save_papers(self, papers: List[Dict]) -> bool:
        """批量保存或更新论文信息到 MongoDB"""
        if not papers:
            logger.warning("No papers to save.")
            return False

        try:
            with MongoClient(self.uri) as client:
                db = client[self.db_name]
                collection = db[self.collection_name]

                # 使用 upsert 实现插入或更新
                for paper in papers:
                    collection.update_one(
                        {"id": paper["id"]},
                        {"$set": paper},
                        upsert=True
                    )
                logger.info(f"Saved or updated {len(papers)} papers in MongoDB.")
                return True
        except Exception as e:
            logger.error(f"Error saving papers to MongoDB: {e}")
            return False

    def update_paper_with_pdf_info(self, paper_id: str, pdf_file_id: Optional[str] = None,
                                   local_pdf_path: Optional[str] = None) -> bool:
        """更新论文记录中的 PDF 路径或 GridFS ID"""
        try:
            with MongoClient(self.uri) as client:
                db = client[self.db_name]
                collection = db[self.collection_name]

                update_data = {}
                if pdf_file_id:
                    update_data["pdf_file_id"] = str(pdf_file_id)
                if local_pdf_path:
                    update_data["local_pdf_path"] = local_pdf_path

                result = collection.update_one(
                    {"id": paper_id},
                    {"$set": update_data}
                )

                if result.modified_count > 0:
                    logger.info(f"Updated PDF info for paper {paper_id}")
                return True
        except Exception as e:
            logger.error(f"Error updating PDF info for paper {paper_id}: {e}")
            return False

    def save_analysis_results(self, results: List[Dict]):
        """保存分析结果到MongoDB"""
        with self.get_client() as client:
            db = client[self.db_name]
            collection = db["analysis_results"]

            for result in results:
                collection.update_one(
                    {"keyword": result["keyword"]},
                    {"$set": result},
                    upsert=True
                )